#include "GeologicalHazard.h"


GeologicalHazardSeg::GeologicalHazardSeg(const std::string &_onnx_path, unsigned int _num_threads,int height, int width) :
    log_id(_onnx_path.data()), num_threads(_num_threads)
{
#ifdef _WIN32
    int ModelPathSize = MultiByteToWideChar(CP_UTF8, 0, _onnx_path.c_str(), static_cast<int>(_onnx_path.length()), nullptr, 0);
    wchar_t* wide_cstr = new wchar_t[ModelPathSize + 1];
    MultiByteToWideChar(CP_UTF8, 0, _onnx_path.c_str(), static_cast<int>(_onnx_path.length()), wide_cstr, ModelPathSize);
    wide_cstr[ModelPathSize] = L'\0';
    const wchar_t* modelPath = wide_cstr;
#else
    const char* modelPath = _onnx_path.c_str();
#endif // _WIN32
  ort_env = Ort::Env(ORT_LOGGING_LEVEL_ERROR, log_id);
  // 0. session options
  Ort::SessionOptions session_options;
  session_options.SetIntraOpNumThreads(num_threads);
  session_options.SetGraphOptimizationLevel(
      GraphOptimizationLevel::ORT_ENABLE_EXTENDED);
  session_options.SetLogSeverityLevel(4);
  // 1. session
  // GPU Compatibility.
#ifdef USE_CUDA
  OrtSessionOptionsAppendExecutionProvider_CUDA(session_options, 0); // C API stable.
#endif
  dynamic_input_height=height; 
  dynamic_input_width=width;
  ort_session = new Ort::Session(ort_env, modelPath, session_options);
  Ort::AllocatorWithDefaultOptions allocator;
  // 2. input name & input dims
  input_node_names.resize(num_inputs); // num_inputs=1
  input_node_names_.resize(num_inputs); // num_inputs=1
  input_node_names_[0] = OrtCompatiableGetInputName(0, allocator, ort_session);
  input_node_names[0] = input_node_names_[0].data();
  // 3. initial input node dims.
  dynamic_input_node_dims.push_back({1, 3, dynamic_input_height, dynamic_input_width});
  dynamic_input_tensor_size = 1 * 3 * dynamic_input_height * dynamic_input_width;
  dynamic_input_values_handler.resize(dynamic_input_tensor_size);
  // 4. output names & output dimms
  num_outputs = ort_session->GetOutputCount();
  output_node_names.resize(num_outputs);
  output_node_names_.resize(num_outputs);
  for (unsigned int i = 0; i < num_outputs; ++i) {
    output_node_names_[i] = OrtCompatiableGetOutputName(i, allocator, ort_session);
    output_node_names[i] = output_node_names_[i].data();
  }
#if LITEORT_DEBUG
  this->print_debug_string();
#endif
}

GeologicalHazardSeg::~GeologicalHazardSeg()
{
  if (ort_session)
    delete ort_session;
  ort_session = nullptr;
}

void GeologicalHazardSeg::print_debug_string()
{
  //std::cout << "LITEORT_DEBUG LogId: " << onnx_path << "\n";
  std::cout << "=============== Inputs ==============\n";
  for (unsigned int i = 0; i < num_inputs; ++i)
    for (unsigned int j = 0; j < dynamic_input_node_dims.at(i).size(); ++j)
      std::cout << "Dynamic Input: " << i << " Name: "
                << input_node_names.at(i) << " Init Dim: " << j << " :"
                << dynamic_input_node_dims.at(i).at(j) << std::endl;
  std::cout << "=============== Outputs ==============\n";
  for (unsigned int i = 0; i < num_outputs; ++i)
    std::cout << "Dynamic Output " << i << ": " << output_node_names[i] << std::endl;
}

Ort::Value GeologicalHazardSeg::transform(const cv::Mat &mat)
{
  cv::Mat canvas;
  const unsigned int img_height = mat.rows;
  const unsigned int img_width = mat.cols;

  // update dynamic input dims
  dynamic_input_height = img_height;
  dynamic_input_width = img_width;
  dynamic_input_node_dims.at(0).at(2) = dynamic_input_height;
  dynamic_input_node_dims.at(0).at(3) = dynamic_input_width;
  dynamic_input_tensor_size = 1 * 3 * dynamic_input_height * dynamic_input_width;
  dynamic_input_values_handler.resize(dynamic_input_tensor_size);

  cv::cvtColor(mat, canvas, cv::COLOR_BGR2RGB);
  canvas.convertTo(canvas, CV_32FC3, 1.f / 255.f, 0.f); // (0.,1.)

  normalize_inplace(canvas, mean_vals, scale_vals);

  return create_tensor(
      canvas, dynamic_input_node_dims.at(0), memory_info_handler,
      dynamic_input_values_handler,CHW
  );

}

void GeologicalHazardSeg::detect(const cv::Mat &mat, SegmentContent &content)
{
  if (mat.empty()) return;
  // 1. make input tensor
  Ort::Value input_tensor = this->transform(mat);
  // 2. inference out(scores) (1,21=1+20,h,w) pixel to pixel
  auto output_tensors = ort_session->Run(
      Ort::RunOptions{nullptr}, input_node_names.data(),
      &input_tensor, num_inputs, output_node_names.data(),
      num_outputs
  );
  // 3. post process.
  Ort::Value &scores = output_tensors.at(0); // (1,2,h,w)
  auto scores_dims = scores.GetTypeInfo().GetTensorTypeAndShapeInfo().GetShape();
  const unsigned int output_classes = scores_dims.at(1);
  const unsigned int output_height = scores_dims.at(2);
  const unsigned int output_width = scores_dims.at(3);

  // time cost!
  content.names_map.clear();
  content.class_mat = cv::Mat(output_height, output_width, CV_8UC1, cv::Scalar(0));
  content.color_mat = mat.clone();

  for (unsigned int i = 0; i < output_height; ++i)
  {

    uchar *p_class = content.class_mat.ptr<uchar>(i);
    cv::Vec3b *p_color = content.color_mat.ptr<cv::Vec3b>(i);

    for (unsigned int j = 0; j < output_width; ++j)
    {
      // argmax
      unsigned int max_label = 0;
      float max_conf = scores.At<float>({0, 0, i, j});

      for (unsigned int l = 0; l < output_classes; ++l)
      {
        float conf = scores.At<float>({0, l, i, j});
        if (conf > max_conf)
        {
          max_conf = conf;
          max_label = l;
        }
      }

      if (max_label == 0) continue;

      // assign label for pixel(i,j)
      p_class[j] = cv::saturate_cast<uchar>(max_label);
      // assign color for detected class at pixel(i,j).
      p_color[j][0] = cv::saturate_cast<uchar>((max_label % 10) * 20);
      p_color[j][1] = cv::saturate_cast<uchar>((max_label % 5) * 40);
      p_color[j][2] = cv::saturate_cast<uchar>((max_label % 10) * 20 );
      // assign names map
      content.names_map[max_label] = class_names[max_label - 1]; // max_label >= 1
    }

  }

  content.flag = true;
}
std::string GeologicalHazardSeg::OrtCompatiableGetInputName(size_t index, OrtAllocator* allocator,
    Ort::Session* ort_session) {
#if ORT_API_VERSION >= 14
    return std::string(ort_session->GetInputNameAllocated(index, allocator).get());
#else  
    return std::string(ort_session->GetInputName(index, allocator));
#endif
}

std::string GeologicalHazardSeg::OrtCompatiableGetOutputName(size_t index, OrtAllocator* allocator,
    Ort::Session* ort_session) {
#if ORT_API_VERSION >= 14
    return std::string(ort_session->GetOutputNameAllocated(index, allocator).get());
#else  
    return std::string(ort_session->GetOutputName(index, allocator));
#endif
}
Ort::Value GeologicalHazardSeg::create_tensor(const cv::Mat& mat,
    const std::vector<int64_t>& tensor_dims,
    const Ort::MemoryInfo& memory_info_handler,
    std::vector<float>& tensor_value_handler,
    unsigned int data_format)
{
    const unsigned int rows = mat.rows;
    const unsigned int cols = mat.cols;
    const unsigned int channels = mat.channels();

    cv::Mat mat_ref;
    if (mat.type() != CV_32FC(channels)) mat.convertTo(mat_ref, CV_32FC(channels));
    else mat_ref = mat; // reference only. zero-time cost. support 1/2/3/... channels

    if (tensor_dims.size() != 4) throw std::runtime_error("dims mismatch.");
    if (tensor_dims.at(0) != 1) throw std::runtime_error("batch != 1");

    // CXHXW
    if (data_format == CHW)
    {

        const unsigned int target_height = tensor_dims.at(2);
        const unsigned int target_width = tensor_dims.at(3);
        const unsigned int target_channel = tensor_dims.at(1);
        const unsigned int target_tensor_size = target_channel * target_height * target_width;
        if (target_channel != channels) throw std::runtime_error("channel mismatch.");

        tensor_value_handler.resize(target_tensor_size);

        cv::Mat resize_mat_ref;
        if (target_height != rows || target_width != cols)
            cv::resize(mat_ref, resize_mat_ref, cv::Size(target_width, target_height));
        else resize_mat_ref = mat_ref; // reference only. zero-time cost.

        std::vector<cv::Mat> mat_channels;
        cv::split(resize_mat_ref, mat_channels);
        // CXHXW
        for (unsigned int i = 0; i < channels; ++i)
            std::memcpy(tensor_value_handler.data() + i * (target_height * target_width),
                mat_channels.at(i).data, target_height * target_width * sizeof(float));

        return Ort::Value::CreateTensor<float>(memory_info_handler, tensor_value_handler.data(),
            target_tensor_size, tensor_dims.data(),
            tensor_dims.size());
    }

    // HXWXC
    const unsigned int target_height = tensor_dims.at(1);
    const unsigned int target_width = tensor_dims.at(2);
    const unsigned int target_channel = tensor_dims.at(3);
    const unsigned int target_tensor_size = target_channel * target_height * target_width;
    if (target_channel != channels) throw std::runtime_error("channel mismatch!");
    tensor_value_handler.resize(target_tensor_size);

    cv::Mat resize_mat_ref;
    if (target_height != rows || target_width != cols)
        cv::resize(mat_ref, resize_mat_ref, cv::Size(target_width, target_height));
    else resize_mat_ref = mat_ref; // reference only. zero-time cost.

    std::memcpy(tensor_value_handler.data(), resize_mat_ref.data, target_tensor_size * sizeof(float));

    return Ort::Value::CreateTensor<float>(memory_info_handler, tensor_value_handler.data(),
        target_tensor_size, tensor_dims.data(),
        tensor_dims.size());
}

cv::Mat GeologicalHazardSeg::normalize(const cv::Mat& mat, float mean, float scale)
{
    cv::Mat matf;
    if (mat.type() != CV_32FC3) mat.convertTo(matf, CV_32FC3);
    else matf = mat; // reference
    return (matf - mean) * scale;
}

cv::Mat GeologicalHazardSeg::normalize(const cv::Mat& mat, const float* mean, const float* scale)
{
    cv::Mat mat_copy;
    if (mat.type() != CV_32FC3) mat.convertTo(mat_copy, CV_32FC3);
    else mat_copy = mat.clone();
    for (unsigned int i = 0; i < mat_copy.rows; ++i)
    {
        cv::Vec3f* p = mat_copy.ptr<cv::Vec3f>(i);
        for (unsigned int j = 0; j < mat_copy.cols; ++j)
        {
            p[j][0] = (p[j][0] - mean[0]) * scale[0];
            p[j][1] = (p[j][1] - mean[1]) * scale[1];
            p[j][2] = (p[j][2] - mean[2]) * scale[2];
        }
    }
    return mat_copy;
}

void GeologicalHazardSeg::normalize(const cv::Mat& inmat, cv::Mat& outmat,
    float mean, float scale)
{
    outmat = normalize(inmat, mean, scale);
}

void GeologicalHazardSeg::normalize_inplace(cv::Mat& mat_inplace, float mean, float scale)
{
    if (mat_inplace.type() != CV_32FC3) mat_inplace.convertTo(mat_inplace, CV_32FC3);
    normalize(mat_inplace, mat_inplace, mean, scale);
}

void GeologicalHazardSeg::normalize_inplace(cv::Mat& mat_inplace, const float* mean, const float* scale)
{
    if (mat_inplace.type() != CV_32FC3) mat_inplace.convertTo(mat_inplace, CV_32FC3);
    for (unsigned int i = 0; i < mat_inplace.rows; ++i)
    {
        cv::Vec3f* p = mat_inplace.ptr<cv::Vec3f>(i);
        for (unsigned int j = 0; j < mat_inplace.cols; ++j)
        {
            p[j][0] = (p[j][0] - mean[0]) * scale[0];
            p[j][1] = (p[j][1] - mean[1]) * scale[1];
            p[j][2] = (p[j][2] - mean[2]) * scale[2];
        }
    }
}
